#include "kernel_operator.h"
#include "op_common.h"
#include <cstdint>
#include <type_traits>

using namespace AscendC;

template <typename _DT_X, typename _DT_Y> class Kernel {
public:
  TPipe pipe;
  DefInTensor(X);
  DefOutTensor(Y);
  int64_t b = 0;
  int64_t w = 0;
  int64_t h = 0;
  int64_t diagonal = 0;
  int64_t wxh = 0;
  int64_t tileLength = 0;

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, T &tiling_data) {
    b = tiling_data.b;
    w = tiling_data.w;
    h = tiling_data.h;
    diagonal = tiling_data.diagonal;
    wxh = w * h;
    tileLength = tiling_data.tileLength;
    auto size = b * w * h;
    // 切分global
    GSetBuffer(X, x, 0, size);
    GSetBuffer(Y, y, 0, size);
    // 初始化队列
    InitQueueSimple(X, tileLength);
    InitQueueSimple(Y, tileLength);
  }

  __aicore__ inline void Process() {
    for (int32_t bi = 0; bi < b; ++bi) {
      for (int32_t i = 0; i < w; ++i) {
        auto seg_index = diagonal + i + 1;
        for (int32_t j = 0; j < h; ++j) {
          auto index = bi * wxh + i * h + j;
          if (j < seg_index) {
            GTensorName(Y).SetValue(index, GTensorName(X).GetValue(index));
          } else {
            GTensorName(Y).SetValue(index, TypeOf(Y)(0));
          }
        }
      }
    }
  }

  __aicore__ inline void DoTrilOnce(int64_t elementCount, int64_t dataOffset,
                                    uint32_t calcCount, int64_t doCount = 1,
                                    int64_t yStep = 0, int64_t dataStep = 0) {
    uint16_t repeatTimes = CEIL_DIV(elementCount, elementsPerRepeat<_DT_X>());
    if (elementCount <= 0) {
      QueAllocSimple(Y);
      if constexpr (std::is_same_v<_DT_X, half>) {
        MY_DUP_INT16(LTensor(Y), zeroHalfInt, tileLength);
      } else if constexpr (std::is_same_v<_DT_X, float>) {
        MY_DUP_INT32(LTensor(Y), zeroFloatInt, tileLength);
      }
      EnQue(Y);
    } else {
      {
        QueAllocSimple(X);
        auto copyCount = ALIGN_TO(elementCount, elementsPerBlock<_DT_X>());
        for (auto i = 0; i < doCount; ++i) {
          DataCopy(LTensor(X)[i * yStep], GTensor(X)[dataOffset + i * dataStep],
                   copyCount);
        }
        EnQue(X);
      }
      {
        DeQueSimple(X);
        QueAllocSimple(Y);
        // print_tensor("X", LTensor(X), 16, "%f");
        // Duplicate(LTensor(Y), _DT_X(0), calcCount);
        if constexpr (std::is_same_v<_DT_X, half>) {
          MY_DUP_INT16(LTensor(Y), zeroHalfInt, tileLength);
        } else if constexpr (std::is_same_v<_DT_X, float>) {
          MY_DUP_INT32(LTensor(Y), zeroFloatInt, tileLength);
        }
        for (auto i = 0; i < doCount; ++i) {
          uint64_t counter = 0;
          GatherMask(LTensor(Y)[i * yStep], LTensor(X)[i * yStep], 7, false, 0,
                     {1, repeatTimes, 8, 8}, counter);
          // DataCopy(LTensor(Y), LTensor(X), copyCount);
          for (auto j = elementCount; j < counter; ++j) {
            LTensor(Y)(i * yStep + j) = _DT_X(0);
          }
        }
        EnQue(Y);
        QueFree(X);
      }
    }
    copyOut(doCount, calcCount, dataOffset, dataStep, yStep);
  }
  __aicore__ inline void copyOut(int64_t doCount, int64_t calcCount,
                                 int64_t dataOffset, int64_t dataStep,
                                 int64_t yStep) {
    DeQueSimple(Y);
    // print_tensor("Y", LTensor(Y), 16, "%f");
    for (auto i = 0; i < doCount; ++i) {
      for (auto j = 0; j < calcCount; ++j) {
        GTensor(Y).SetValue(dataOffset + i * dataStep + j,
                            LTensor(Y).GetValue(i * yStep + j));
      }
      // if (calcCount < elementsPerRepeat<_DT_X>()) {
      //   for (auto j = 0; j < calcCount; ++j) {
      //     GTensor(Y).SetValue(dataOffset + i * dataStep + j,
      //                         LTensor(Y).GetValue(i * yStep + j));
      //   }
      // } else {
      //   // myDataCopy(GTensor(Y)[dataOffset], LTensor(Y), calcCount);
      //   DataCopy(GTensor(Y)[dataOffset + i * dataStep], LTensor(Y)[i *
      //   yStep],
      //            ALIGN_TO(calcCount, elementsPerBlock<_DT_X>()));
      //   DataCopyPad(GTensor(Y)[dataOffset + i * dataStep],
      //               LTensor(Y)[i * yStep], {1, 256, 0, 0});
      // }
    }
    QueFree(Y);
  }

  __aicore__ inline int64_t getElementCount(int64_t wi) {
    return GET_MIN(diagonal + w + 1, w);
  }

  __aicore__ inline void DoTril(int64_t loopCount, int64_t bi, int64_t hCount,
                                int64_t yStep, int64_t finalLength) {
    for (int32_t i = 0; i < w; ++i) {
      auto elementCount = getElementCount(i);
      auto dataOffset = bi * wxh + i * h;
      auto calcCount = GET_MIN(h, tileLength);
      for (auto k = 0; k < loopCount; ++k) {
        auto loopOffset = k * tileLength;
        DoTrilOnce(elementCount < loopOffset ? 0 : elementCount,
                   dataOffset + loopOffset, calcCount, hCount, yStep, wxh);
      }
      if (finalLength > 0) {
        DoTrilOnce(elementCount < loopCount * tileLength ? 0 : elementCount,
                   dataOffset + loopCount * tileLength, finalLength, hCount,
                   yStep, wxh);
      }
    }
  }
  __aicore__ inline void DoTrilMutiWOnce(int64_t dataOffset,
                                         int64_t elementCount, int64_t hCount,
                                         int64_t alignH) {
    {
      QueAllocSimple(X);
      auto tempElementCount = elementCount;
      for (auto i = 0; i < hCount; ++i) {
        auto copyCount = ALIGN_TO(elementCount, elementsPerBlock<_DT_X>());
        DataCopy(LTensor(X)[i * alignH], GTensor(X)[dataOffset + i * h],
                 copyCount);
        tempElementCount++;
      }
      EnQue(X);
    }
    {
      DeQueSimple(X);
      QueAllocSimple(Y);
      // print_tensor("X", LTensor(X), 16, "%f");
      // Duplicate(LTensor(Y), _DT_X(0), calcCount);
      if constexpr (std::is_same_v<_DT_X, half>) {
        MY_DUP_INT16(LTensor(Y), zeroHalfInt, tileLength);
      } else if constexpr (std::is_same_v<_DT_X, float>) {
        MY_DUP_INT32(LTensor(Y), zeroFloatInt, tileLength);
      }
      auto tempElementCount = elementCount;
      for (auto i = 0; i < hCount; ++i) {
        if (tempElementCount <= 0) {
          tempElementCount++;
          continue;
        }
        uint64_t counter = 0;
        uint16_t repeatTimes =
            CEIL_DIV(tempElementCount, elementsPerRepeat<_DT_X>());
        GatherMask(LTensor(Y)[i * alignH], LTensor(X)[i * alignH], 7, false, 0,
                   {1, repeatTimes, 8, 8}, counter);
        // DataCopy(LTensor(Y), LTensor(X), copyCount);
        for (auto j = tempElementCount; j < counter; ++j) {
          LTensor(Y)(i * alignH + j) = _DT_X(0);
        }
        // print("xxxxxxxxxxxxxxxxxxx\n");
        // print_tensor("X1", LTensor(X), 3, "%f");
        // print_tensor("Y1", LTensor(Y), 3, "%f");
        // print_tensor("X2", LTensor(X)[elementsPerBlock<_DT_X>()], 3, "%f");
        // print_tensor("Y2", LTensor(Y)[elementsPerBlock<_DT_X>()], 3, "%f");
        // print_tensor("X3", LTensor(X)[elementsPerBlock<_DT_X>() * 2], 3,
        //              "%f");
        // print_tensor("Y3", LTensor(Y)[elementsPerBlock<_DT_X>() * 2], 3,
        //              "%f");
        // print("xxxxxxxxxxxxxxxxxxx\n");
        tempElementCount++;
      }
      EnQue(Y);
      QueFree(X);
    }
    copyOut(hCount, h, dataOffset, h, alignH);
  }

  __aicore__ inline void Process2() {
    auto alignH = ALIGN_TO(h, elementsPerBlock<_DT_X>());
    auto hCount = tileLength / alignH;
    // 对于h比较大的情况
    if (hCount <= 0) {
      auto loopCount = h / tileLength;
      auto finalLength = h % tileLength;
      for (int32_t bi = 0; bi < b; ++bi) {
        DoTril(loopCount, bi, 1, tileLength, finalLength);
      }
      return;
    }

    // 对于h比较小， 但是b很大的情况
    if (b >= w) {
      hCount = GET_MIN(hCount, b);
      for (int32_t bi = 0; bi < b; bi += hCount) {
        DoTril(0, bi, hCount, alignH, h);
      }
      auto lastB = b % hCount;
      if (lastB > 0) {
        DoTril(0, b - lastB, lastB, alignH, h);
      }
      return;
    }
    hCount = GET_MIN(hCount, b);
    auto copyCount = alignToBlock<_DT_X>(hCount * h);
    for (int32_t bi = 0; bi < b; ++bi) {
      for (auto wi = 0; wi < w; wi += hCount) {
        DoTrilMutiWOnce(bi * wxh + wi * h, getElementCount(wi), hCount, alignH);
      }
      auto lastW = w % hCount;
      if (lastW > 0) {
        auto dataOffset = bi * wxh + (w - lastW) * h;
        DoTrilMutiWOnce(dataOffset, getElementCount(w - lastW), lastW, alignH);
      }
    }
  }
};

template <typename _DT_X, typename _DT_Y> class Kernel2 {
public:
  TPipe pipe;
  DefInTensor(X);
  DefOutTensor(Y);
  int64_t b = 0;
  int64_t w = 0;
  int64_t h = 0;
  int64_t diagonal = 0;
  int64_t wxh = 0;
  int64_t tileLength = 0;
  int64_t bi = 0;
  int64_t wi = 0;
  int64_t hi = 0;
  int64_t size = 0;

public:
  __aicore__ inline Kernel2() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, T &tiling_data) {
    b = tiling_data.b;
    w = tiling_data.w;
    h = tiling_data.h;
    diagonal = tiling_data.diagonal;
    wxh = w * h;
    tileLength = tiling_data.tileLength;
    size = b * w * h;
    // 切分global
    GSetBuffer(X, x, 0, size);
    GSetBuffer(Y, y, 0, size);
    // 初始化队列
    InitQueueSimple(X, tileLength);
    InitQueueSimple(Y, tileLength);
  }

  __aicore__ inline bool increaseHi(int64_t v) {
    hi += v;
    if (hi >= h) {
      wi += hi / h;
      hi = hi % h;
      if (wi >= w) {
        bi += wi / w;
        wi = wi % w;
        if (bi >= b) {
          return true;
        }
      }
    }
    return false;
  }

  __aicore__ inline int64_t getZeroCount() {
    return GET_MIN(diagonal + wi + 1 - hi, h - hi);
  }

#ifdef ASCENDC_DUMP
  template <typename T> __aicore__ inline void print_state(T desc) {
    print("xxxxxxxxxxxxxxx%sxxxxxxxxxxxx\n", desc);
    auto dataCount = getZeroCount();
    print("bi = %d\n", bi);
    print("wi = %d\n", wi);
    print("hi = %d\n", hi);
    print("dataCount = %d\n", dataCount);
  }
#else
#define print_state(desc)
#endif

  __aicore__ inline void Process() {
    auto loopCount = CEIL_DIV(size, tileLength);
    for (auto i = 0; i < loopCount; ++i) {
      auto calcCount = tileLength;
      auto copyCount = tileLength;
      if (i == loopCount - 1) {
        calcCount = size % tileLength;
        copyCount = alignToBlock<_DT_X>(calcCount);
      }
      if (diagonal <= -w) {
        QueAllocSimple(Y);
        Muls(LTensor(Y), LTensor(Y), _DT_X(0), copyCount);
        EnQue(Y);
      } else {
        { EnQueGlobal2Local(X, i * tileLength, copyCount); }
        {
          DeQueSimple(X);
          QueAllocSimple(Y);
          Muls(LTensor(Y), LTensor(Y), _DT_X(0), copyCount);
          if (diagonal >= w - 1 && diagonal >= h - 1) {
            Add(LTensor(Y), LTensor(X), LTensor(Y), calcCount);
          } else {
            auto tensorIdx = 0;
            auto extraIdx = 0;
            while (tensorIdx < calcCount) {
              auto dataCount = getZeroCount();
              // 刚好可以占满一个repeat
              if (dataCount >= elementsPerRepeat<_DT_X>()) {
                auto count = GET_MIN(dataCount, copyCount - tensorIdx) /
                             elementsPerRepeat<_DT_X>() *
                             elementsPerRepeat<_DT_X>();
                increaseHi(count);
                Add(LTensor(Y)[tensorIdx], LTensor(Y)[tensorIdx],
                    LTensor(X)[tensorIdx], count);
                tensorIdx += count;
                continue;
              }
              // 构建mask
              uint64_t mask[2] = {0, 0};
              int32_t maskFlagIdx = 0;
              while (maskFlagIdx < elementsPerRepeat<_DT_X>()) {
                // 对于只有0的情况
                if (dataCount <= 0) {
                  // 统计该行还有多少个0
                  auto zeroCount = h - hi;
                  auto newMaskFlagIdx = maskFlagIdx + zeroCount;
                  // 如果超出mask承载，则表示该次repeat全是0，不用Add,直接结束
                  if (newMaskFlagIdx >= elementsPerRepeat<_DT_X>()) {
                    auto increaseCount =
                        elementsPerRepeat<_DT_X>() - maskFlagIdx;
                    auto lastZeroCount = zeroCount - increaseCount;
                    extraIdx =
                        GET_MIN(lastZeroCount, copyCount - tensorIdx -
                                                   elementsPerRepeat<_DT_X>()) /
                        elementsPerRepeat<_DT_X>() * elementsPerRepeat<_DT_X>();
                    increaseHi(increaseCount + extraIdx);
                    break;
                  }
                  maskFlagIdx = newMaskFlagIdx;
                  // 判断是否需要直接结束， 因为都是0了，不用再加了
                  if (increaseHi(zeroCount)) {
                    break;
                  }
                  // 重新获取数据个数
                  dataCount = getZeroCount();
                  continue;
                }
                dataCount = GET_MIN(dataCount,
                                    elementsPerRepeat<_DT_X>() - maskFlagIdx);
                MASK_OR_ONE(_DT_X, mask, maskFlagIdx, dataCount);
                if (maskFlagIdx >= elementsPerRepeat<_DT_X>()) {
                  increaseHi(dataCount);
                  goto maskCalc;
                }
                // 判断是否已经没有数据了
                if (increaseHi(dataCount)) {
                  goto maskCalc;
                }

                // 重新获取数据个数
                dataCount = getZeroCount();
              }
            maskCalc:
              Add(LTensor(Y)[tensorIdx], LTensor(Y)[tensorIdx],
                  LTensor(X)[tensorIdx], mask, 1, {1, 1, 1, 8, 8, 8});
              tensorIdx += elementsPerRepeat<_DT_X>() + extraIdx;
              extraIdx = 0;
            }
          }
          EnQue(Y);
          QueFree(X);
        }
      }
      { DeQueLocal2Global(Y, i * tileLength, copyCount); }
    }
  }
};

extern "C" __global__ __aicore__ void tril(GM_ADDR x, GM_ADDR y,
                                           GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  // Kernel<DTYPE_X, DTYPE_Y> op;
  // op.Init(x, y, tiling_data);
  // // op.Process();
  // op.Process2();

  Kernel2<DTYPE_X, DTYPE_Y> op;
  op.Init(x, y, tiling_data);
  op.Process();
}